'''
Created on Feb 5, 2013

@author: nzhao
'''
import numpy as np
import scipy.sparse
from scipy import linalg
from itertools import combinations

from MatViewer.MatViewer import MatViewer

from traits.api import Instance, Array, Bool
from math_qm.QMobj import QMobject
from math_qm.QMoperator import QMoperator
from phy.interaction.SingleSpinInteraction import SingleSpinInteraction
from phy.interaction.MultiSpinInteraction import MultiSpinInteraction


class MultiLevel_EigenSystem(QMoperator):
    col_list = Array()
    row_list = Array()
    val_list = Array()
    
    eig_Vals = Array()
    eig_Vecs = Array()
    
    mat_sparse = Instance(scipy.sparse.coo_matrix)
    has_mat    = Bool(False)
    isHermite  = Bool(False)

    def __init__(self, nbody, dim_list):
        QMobject.__init__(self, nbody, dim_list)
        
        self.col_list = np.array([])
        self.row_list = np.array([])
        self.val_list = np.array([])
        
    def calculate_mat(self, n_coupling_lst=[1,2], is2dense=True):
        idx_list = range(self.nbody)
        for n_coupling in n_coupling_lst:
            idx_sel = list(combinations(idx_list, n_coupling))
            for idx in idx_sel:
                self.cal_elements(idx)
        
        self.mat_sparse = scipy.sparse.coo_matrix((self.val_list, (np.array(self.row_list), np.array(self.col_list))), shape=(self.dim, self.dim))
        
        if is2dense:
            self.has_mat = True
            self.mat = np.matrix( self.mat_sparse.todense() )
            self.isHermite = np.array_equal( self.mat.H, self.mat)
            if not self.isHermite:
                print 'error: matrix is not hermitian.'
        
    def cal_elements(self, sel_idx):
        print self.dim_list
        print sel_idx
        print 
        sel_dim_lst = [self.dim_list[i]  for i in sel_idx]
        sel_dim     = np.prod(sel_dim_lst)
        states      = self._generate_basis(sel_dim_lst, sel_dim)
        amplitudes = self.scattering_amplitude(sel_idx, sel_dim, states)
        
        col_ = np.array([])
        row_ = np.array([])
        val_ = np.array([])
        
        for x in range(sel_dim):
            for y in range(sel_dim):
                amp = amplitudes[x][y]
                
                row_indices = [i for i in range(self.dim) if self._matchQ(self.basis_list[i], sel_idx, states[x])]
                col_indices = [i for i in range(self.dim) if self._matchQ(self.basis_list[i], sel_idx, states[y])]
                
                len1 = len(row_indices)
                 
                if np.abs(amp) > 0.0:
                    val_ = np.append(val_, amp * np.ones(len1) )
                    col_ = np.append(col_, np.array(col_indices))
                    row_ = np.append(row_, np.array(row_indices))

        self.col_list = np.append(self.col_list, col_)
        self.row_list = np.append(self.row_list, row_)
        self.val_list = np.append(self.val_list, val_)
        
    def scattering_amplitude(self, idx, dim, states):
        return np.zeros((dim, dim))
    
    def _matchQ(self, basis, sel_idx, state):
        basis_sel = [basis[i] for i in sel_idx]
        return basis_sel == state
    
    def eigen_system(self):
        if self.has_mat and self.isHermite:
            val, vec = linalg.eigh(self.mat)
            self.eig_Vals = val
            self.eig_Vecs = vec.T
        else:
            print "Has dense mat = %d, HermitianQ = %d" %(self.has_mat, self.isHermite) 
            
        
class SingleSpinEigen(MultiLevel_EigenSystem):
    spin = Instance(SingleSpinInteraction)
    
class Eigen_NSpinHalf(MultiLevel_EigenSystem):
    spins = Instance(MultiSpinInteraction)

    def __init__(self, spins, is_auto_update=True):
        try:
            self.spins = spins
        except:
            print 'interacting spins are needed.'
            
        self.nbody = self.spins.nspin
        MultiLevel_EigenSystem.__init__(self, self.nbody, [spin.dim for spin in self.spins.spin_list])
        
        if is_auto_update:
            self.calculate_mat()
    
    def scattering_amplitude(self, idx, dim, states):
        len_states = len(idx)
        
        if len_states == 1:
            bx, by, bz = self.spins.spin_list[ idx[0] ].get_local_field()
            res = 0.5 * np.array([[bz,               bx - 1.j * by],
                                  [bx + 1.j * by,   -bz           ]
                                  ])
        elif len_states == 2:
            i, j = idx
            
            dip = self.spins.dipolar_tensor(i, j)
            dip_zp = dip.zx + 1.j * dip.zy
            dip_zm = dip.zx - 1.j * dip.zy
            dip_ss = dip.xx + dip.yy
            dip_s_s= dip.xx - dip.yy
            dip_cc = dip.xy + dip.yx
            
            res = 0.25*np.array(
                            [[  dip.zz,                 - dip_zp,     - dip_zp,   dip_s_s + 1.j * dip_cc ],
                             [- dip_zm,                 - dip.zz,       dip_ss,   dip_zp                 ],
                             [- dip_zm,                   dip_ss,     - dip.zz,   dip_zp                 ],
                             [  dip_s_s -1.j * dip_cc,    dip_zm,       dip_zm,   dip.zz                 ],
                             ]
                            )
        else:
            res = np.zeros((dim, dim))
        
        return res

if __name__=='__main__':
    from math_qm.vector import vector
    
    applied_b = vector(xyz=[0., 0., 100.])
    
    s1=SingleSpinInteraction(element='C13', spin_type='Bath', x=0.0, y=0.0, z=0.0, b=applied_b)
    s2=SingleSpinInteraction(element='C13', spin_type='Bath', x=0.0, y=0.0, z=1.54, b=applied_b)
    s3=SingleSpinInteraction(element='C13', spin_type='Bath', x=1.54, y=0.0, z=0.0, b=applied_b)
    s4=SingleSpinInteraction(element='C13', spin_type='Bath', x=1.0, y=3.0, z=6.0, b=applied_b)
    s5=SingleSpinInteraction(element='C13', spin_type='Bath', x=1.0, y=3.0, z=16.0, b=applied_b)
    s6=SingleSpinInteraction(element='C13', spin_type='Bath', x=1.0, y=3.0, z=26.0, b=applied_b)
    
    spins = MultiSpinInteraction(spins=[s1,s4])#, s3, s4, s5, s6])
    spin_eigs = Eigen_NSpinHalf(spins)

    a = spin_eigs.mat

    view_a = MatViewer(a/2.0/np.pi)
    view_a.edit_traits()
    
    spin_eigs.eigen_system()
    print spin_eigs.eig_Vals
    eig_a = spin_eigs.eig_Vecs
    view_eig_a = MatViewer(eig_a)
    view_eig_a.configure_traits()

































    
